import tiktoken
import subprocess
import sys

enc = tiktoken.get_encoding("gpt2")
if len(sys.argv) == 1:
    text = input("Text to complete: ")
else:
    text = ' '.join(sys.argv[1:])
    print(text)

tokens = [
    str(tok) for tok in enc.encode(text)
]
tokens.insert(0, str(20))

proc = subprocess.Popen(
    ["./cmake-build-release/gpt2", *tokens],
    stdout=subprocess.PIPE,
    text=True
)

while line := proc.stdout.readline():
    token = int(line)
    print(enc.decode([token]), end='', flush=True)
